package com.kfyang.manualgear.generator.impl.jdbc;

import com.kfyang.manualgear.constant.Constants;
import com.kfyang.manualgear.generator.declare.ClassWrapper;
import com.kfyang.manualgear.generator.declare.QueryMethodGenerator;
import com.kfyang.manualgear.struct.Handler;
import com.kfyang.manualgear.struct.ParameterInfo;
import com.kfyang.manualgear.util.GlobalConfigUtils;
import com.kfyang.manualgear.util.StringUtils;
import com.squareup.javapoet.CodeBlock;
import com.squareup.javapoet.FieldSpec;
import com.squareup.javapoet.ParameterizedTypeName;

import javax.lang.model.element.ExecutableElement;
import javax.lang.model.element.Modifier;
import java.util.Set;
import java.util.StringJoiner;

import static com.kfyang.manualgear.constant.Constants.EMPTY;
import static com.kfyang.manualgear.generator.impl.jdbc.constant.ClassConstants.*;
import static com.kfyang.manualgear.generator.impl.jdbc.constant.StatementConstants.*;
import static com.kfyang.manualgear.util.CheckUtils.checkNotNull;

public class QueryMethodGeneratorJdbcImpl extends QueryMethodGenerator {

    public QueryMethodGeneratorJdbcImpl(ExecutableElement element, ClassWrapper classWrapper, boolean useSimple) {
        super(element, classWrapper, useSimple);
    }

    @Override
    protected void addReturn() {
        ParameterizedTypeName rowMapperType = ParameterizedTypeName.get(ROW_MAPPER_CLASS, returnTypeClass);
        if (methodInfo.resultMap != null) { // 处理@ResultMap

            methodBuilder.beginControlFlow("$T rowMapper = " + LAMBDA_START, rowMapperType);
            methodBuilder.addStatement(NEW_ENTITY_STATEMENT, returnTypeClass, returnTypeClass);

            entitySetAttributes();
            methodBuilder.addStatement(RETURN_ENTITY);

            methodBuilder.endControlFlow(EMPTY);
        } else { // 处理@Construct
            String mapperField = methodInfo.construct + Constants.MAPPER;

            methodBuilder.addStatement("$T rowMapper = $L", rowMapperType, mapperField);

            Set<String> constructSet = classWrapper.getConstructSet(returnTypeClass);
            if (!constructSet.contains(methodInfo.construct)) {
                constructSet.add(methodInfo.construct);

                FieldSpec rowMapperSpec = FieldSpec.builder(rowMapperType,
                        mapperField, Modifier.PRIVATE, Modifier.FINAL)
                    .initializer(rowMapperInitializer())
                    .build();
                classWrapper.addField(rowMapperSpec);
            }
        }
        methodBuilder.addStatement(QUERY_RETURN_STATEMENT, returnTypeInfo.className, returnTypeInfo.methodName, returnTypeInfo.parameters);
    }

    /**
     * rowMapper赋值
     * @return rowMapper赋值语句的代码块
     */
    private CodeBlock rowMapperInitializer() {
        CodeBlock.Builder initializer = CodeBlock.builder();
        initializer.beginControlFlow(LAMBDA_START);

        getValueFromRs(initializer);

        StringJoiner parameterList = new StringJoiner(Constants.COMMA);
        for (ParameterInfo parameterInfo : constructorList) {
            parameterList.add(parameterInfo.name);
        }
        initializer.addStatement(RETURN_ENTITY_BY_CONSTRUCT, returnTypeClass, parameterList.toString());
        initializer.endControlFlow();
        return initializer.build();
    }

    /**
     * 给属性赋值
     */
    private void entitySetAttributes() {
        int colPos = 0;
        for (String resultField : methodInfo.resultMap) {
            Handler handler = fieldToHandler.get(resultField);
            checkNotNull(handler, "没有对应的类型处理器");
            methodBuilder.addStatement(ENTITY_SET_STATEMENT,
                StringUtils.toSetMethod(resultField), handler.className, handler.methodName, GlobalConfigUtils.alias[colPos++]);
        }
    }

    /**
     * 从 resultSet 中获取属性值
     * @param initializer 赋值语句代码块构造器
     */
    private void getValueFromRs(CodeBlock.Builder initializer) {
        int colPos = 0;
        for (ParameterInfo parameterInfo : constructorList) {
            initializer.addStatement(GET_VALUE_STATEMENT,
                parameterInfo.type, parameterInfo.name,
                parameterInfo.handler.className, parameterInfo.handler.methodName,
                GlobalConfigUtils.alias[colPos++]);
        }
    }
}
